import matplotlib.pyplot as plt
import matplotlib

X = [1.0, 2.0, 3.0]
Y = [2.0, 4.0, 6.0]

w = 1.0

def forward(x, w):
    return x * w

def cost(X, Y):
    cost = 0
    for x, y in zip(X, Y):
        y_pred = forward(x, w)
        cost += (y_pred - y) ** 2
    return cost / len(X) # 平均损失

def gradient(X, Y):
    grad = 0
    for x, y in zip(X, Y):
        grad += 2 * x * (x * w - y)
    return grad / len(X)

print("预测值（开始训练前）：", 4, forward(4, w))

epochs = 100 # 迭代次数
lr = 0.01 # 学习速率
loss_list = []
epochs_list = []

for epoch in range(epochs):
    cost_val = cost(X, Y)
    grad_val = gradient(X, Y)
    w -= lr * grad_val
    print("Epoch: ", epoch, " w=", "%.2f" % w, " loss=", "%.2f" % cost_val)
    loss_list.append(cost_val)
    epochs_list.append(epoch)

print("预测值（训练结束后）：", "%.2f" % forward(4, w))

# 绘图
matplotlib.use("TkAgg")
plt.plot(epochs_list, loss_list)
plt.xlabel("Epoch")
plt.ylabel("Cost")
plt.show()